import torch
import torch.nn as nn
import torch.optim as optim
import torchattacks
from torchvision import transforms
from torch.utils.data import DataLoader
import os
import matplotlib.pyplot as plt
from sar_cnn import SimpleCNN, MSTAR_Dataset

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 加载 SimpleCNN 模型
output_dir = "./output"
model_save_path = os.path.join(output_dir, "sar_cnn_model.pth")
num_classes = len(os.listdir('./data/mstar/train'))
model = SimpleCNN(num_classes=num_classes).to(device)
model.load_state_dict(torch.load(model_save_path))
model.eval()
print(f"已从 '{model_save_path}' 加载训练好的模型，并切换到评估模式")

# 数据集和数据加载器
transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=3),
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

test_dataset = MSTAR_Dataset(root_dir='./data/mstar/test', transform=transform)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
train_dataset = MSTAR_Dataset(root_dir='./data/mstar/train', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# 定义攻击方法
attacks = {
    'FGSM': torchattacks.FGSM(model, eps=8/255),
    'PGD': torchattacks.PGD(model, eps=8/255, alpha=2/255, steps=40),
    'CW': torchattacks.CW(model, c=1, steps=1000, lr=0.01),
    'DeepFool': torchattacks.DeepFool(model, steps=50),
    'BIM': torchattacks.BIM(model, eps=8/255, alpha=2/255, steps=10)
}

# 计算攻击成功率
def evaluate_attack_success_rate(model, attack, data_loader):
    model.eval()
    total, success = 0, 0
    for images, labels in data_loader:
        images, labels = images.to(device), labels.to(device)
        adv_images = attack(images, labels)
        outputs = model(adv_images)
        _, predicted = outputs.max(1)
        success += (predicted != labels).sum().item()
        total += labels.size(0)
    return 100 * success / total

# 对抗训练前的攻击成功率
before_training_success_rates = {}
print("Attack Success Rates (Before Adversarial Training):")
for name, attack in attacks.items():
    success_rate = evaluate_attack_success_rate(model, attack, test_loader)
    before_training_success_rates[name] = success_rate
    print(f"{name}: {success_rate:.2f}%")

# 对抗训练函数
def adversarial_training(model, train_loader, optimizer, criterion, attack, epochs=5):
    model.train()
    for epoch in range(epochs):
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            adv_images = attack(images, labels)
            optimizer.zero_grad()
            loss = (criterion(model(images), labels) + criterion(model(adv_images), labels)) / 2
            loss.backward()
            optimizer.step()
        print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}")

# 设置训练参数
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# 使用 CW 攻击进行对抗训练
adversarial_training(model, train_loader, optimizer, criterion, attacks['PGD'], epochs=5)

# 对抗训练后的攻击成功率
after_training_success_rates = {}
print("\nAttack Success Rates (After Adversarial Training):")
for name, attack in attacks.items():
    success_rate = evaluate_attack_success_rate(model, attack, test_loader)
    after_training_success_rates[name] = success_rate
    print(f"{name}: {success_rate:.2f}%")

# 可视化攻击成功率的变化
def plot_attack_success_rates(before, after):
    attack_names = list(before.keys())
    before_rates = [before[name] for name in attack_names]
    after_rates = [after[name] for name in attack_names]

    x = range(len(attack_names))  # 横坐标

    fig, ax = plt.subplots(figsize=(10, 6))
    ax.bar(x, before_rates, width=0.4, label='Before Training', align='center')
    ax.bar([i + 0.4 for i in x], after_rates, width=0.4, label='After Training', align='center')

    ax.set_xlabel('Attack Methods')
    ax.set_ylabel('Attack Success Rate (%)')
    ax.set_title('Attack Success Rate Before and After Adversarial Training')
    ax.set_xticks([i + 0.2 for i in x])
    ax.set_xticklabels(attack_names)
    ax.legend()

    plt.tight_layout()
    plt.show()

# 绘制成功率的柱状图
plot_attack_success_rates(before_training_success_rates, after_training_success_rates, save_path="./output/attack_success_rates.png")
